//! Memory allocation testing framework in Rust.
//! This module provides custom implementations of memory allocation functions
//! (`malloc`, `free`, `realloc`, `calloc`, `strdup`) with additional features
//! for testing purposes, such as memory safety checks and debug patterns.

use std::ptr::NonNull;
use std::mem;
use std::slice;

/// Magic number used to identify valid allocated blocks.
const ALLOC_TEST_MAGIC: u32 = 0x72ec82d2;

/// Pattern written to memory after it is freshly allocated.
const MALLOC_PATTERN: u32 = 0xBAADF00D;

/// Pattern written to memory after it is freed.
const FREE_PATTERN: u32 = 0xDEADBEEF;

/// Metadata header for each allocated block.
#[derive(Debug)]
struct BlockHeader {
    magic_number: u32,
    bytes: usize,
}

/// Count of the current number of allocated bytes.
static mut ALLOCATED_BYTES: usize = 0;

/// Limit on the number of allocations that are possible.
static mut ALLOCATION_LIMIT: i32 = -1;

/// Overwrites a block of memory with a repeated pattern.
fn overwrite(ptr: *mut u8, length: usize, pattern: u32) {
    unsafe {
        let byte_ptr = slice::from_raw_parts_mut(ptr, length);
        for (i, b) in byte_ptr.iter_mut().enumerate() {
            let pattern_seq = (i & 3) as u32;
            *b = ((pattern >> (8 * pattern_seq)) & 0xff) as u8;
        }
    }
}

/// Retrieves the block header for a given allocated pointer.
/// # Safety
/// This function assumes that `ptr` is a valid pointer to an allocated block.
unsafe fn get_header(ptr: *mut u8) -> NonNull<BlockHeader> {
    let header_ptr = (ptr as *mut BlockHeader).offset(-1);
    let header = &mut *header_ptr;
    assert_eq!(header.magic_number, ALLOC_TEST_MAGIC, "Invalid magic number");
    NonNull::new(header_ptr).expect("Header pointer should not be null")
}

/// Custom implementation of `malloc`.
/// Allocates a block of memory with the specified size and initializes it with a pattern.
/// # Safety
/// This function uses unsafe Rust to manage raw pointers and memory allocation.
pub unsafe fn alloc_test_malloc(bytes: usize) -> *mut u8 {
    if unsafe { ALLOCATION_LIMIT == 0 } {
        return std::ptr::null_mut();
    }

    let header_size = mem::size_of::<BlockHeader>();
    let total_size = header_size + bytes;
    let header_ptr = std::alloc::alloc(std::alloc::Layout::from_size_align(total_size, 8).unwrap());

    if header_ptr.is_null() {
        return std::ptr::null_mut();
    }

    let header = &mut *(header_ptr as *mut BlockHeader);
    header.magic_number = ALLOC_TEST_MAGIC;
    header.bytes = bytes;

    let ptr = header_ptr.add(header_size);
    overwrite(ptr, bytes, MALLOC_PATTERN);

    unsafe {
        ALLOCATED_BYTES += bytes;
        if ALLOCATION_LIMIT > 0 {
            ALLOCATION_LIMIT -= 1;
        }
    }

    ptr
}

/// Custom implementation of `free`.
/// Frees the allocated block and overwrites it with a pattern to detect misuse.
/// # Safety
/// This function assumes that `ptr` is a valid pointer to an allocated block.
pub unsafe fn alloc_test_free(ptr: *mut u8) {
    if ptr.is_null() {
        return;
    }

    let header = get_header(ptr);
    let block_size = header.as_ref().bytes;
    assert!(unsafe { ALLOCATED_BYTES >= block_size }, "Invalid allocation size");

    overwrite(ptr, block_size, FREE_PATTERN);
    header.as_mut().magic_number = 0;

    let header_ptr = header.as_ptr();
    let header_size = mem::size_of::<BlockHeader>();
    let total_size = header_size + block_size;
    std::alloc::dealloc(
        header_ptr as *mut u8,
        std::alloc::Layout::from_size_align(total_size, 8).unwrap(),
    );

    unsafe {
        ALLOCATED_BYTES -= block_size;
    }
}

/// Custom implementation of `realloc`.
/// Reallocates the memory block to the new size and copies the old data.
/// # Safety
/// This function assumes that `ptr` is a valid pointer to an allocated block.
pub unsafe fn alloc_test_realloc(ptr: *mut u8, bytes: usize) -> *mut u8 {
    let new_ptr = alloc_test_malloc(bytes);
    if new_ptr.is_null() {
        return std::ptr::null_mut();
    }

    if !ptr.is_null() {
        let header = get_header(ptr);
        let bytes_to_copy = std::cmp::min(header.as_ref().bytes, bytes);
        std::ptr::copy_nonoverlapping(ptr, new_ptr, bytes_to_copy);
        alloc_test_free(ptr);
    }

    new_ptr
}

/// Custom implementation of `calloc`.
/// Allocates a block of memory and initializes it to zero.
pub unsafe fn alloc_test_calloc(nmemb: usize, bytes: usize) -> *mut u8 {
    let total_bytes = nmemb * bytes;
    let ptr = alloc_test_malloc(total_bytes);
    if !ptr.is_null() {
        std::ptr::write_bytes(ptr, 0, total_bytes);
    }
    ptr
}

/// Custom implementation of `strdup`.
/// Duplicates the given string into a newly allocated block.
pub unsafe fn alloc_test_strdup(string: &str) -> *mut u8 {
    let len = string.len();
    let ptr = alloc_test_malloc(len + 1);
    if !ptr.is_null() {
        std::ptr::copy_nonoverlapping(string.as_ptr(), ptr, len);
        *ptr.add(len) = 0; // Null-terminate the string
    }
    ptr
}

/// Sets the allocation limit.
pub fn alloc_test_set_limit(alloc_count: i32) {
    unsafe {
        ALLOCATION_LIMIT = alloc_count;
    }
}

/// Retrieves the current number of allocated bytes.
pub fn alloc_test_get_allocated() -> usize {
    unsafe { ALLOCATED_BYTES }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_malloc_free() {
        unsafe {
            let ptr = alloc_test_malloc(10);
            assert!(!ptr.is_null());
            alloc_test_free(ptr);
        }
    }

    #[test]
    fn test_realloc() {
        unsafe {
            let ptr = alloc_test_malloc(10);
            assert!(!ptr.is_null());
            let new_ptr = alloc_test_realloc(ptr, 20);
            assert!(!new_ptr.is_null());
            alloc_test_free(new_ptr);
        }
    }

    #[test]
    fn test_calloc() {
        unsafe {
            let ptr = alloc_test_calloc(5, 4);
            assert!(!ptr.is_null());
            alloc_test_free(ptr);
        }
    }

    #[test]
    fn test_strdup() {
        unsafe {
            let ptr = alloc_test_strdup("hello");
            assert!(!ptr.is_null());
            alloc_test_free(ptr);
        }
    }
}
